package com.edata.bigdata.nebula

import com.vesoft.nebula.connector.NebulaConnectionConfig
import com.vesoft.nebula.connector.ssl.SSLSignType


trait NeConnector {
  var META_ADDRESS = "127.0.0.1:9559"
  var GRAPH_ADDRESS = "127.0.0.1:9779"
  var META_SSL = false
  var GRAPH_SSL = false
  var SSL_SIGN_TYPE = "CA"
  var SSL_CA_PARAM_PEM = "src/main/resources/ssl/casigned.pem"
  var SSL_CA_PARAM_CRT = "/src/main/resources/ssl/casigned.crt"
  var SSL_CA_PARAM_KEY = "src/main/resources/ssl/casigned.key"
  var SSL_SELF_PARAM_PEM = "src/main/resources/ssl/selfsigned.pem"
  var SSL_SELF_PARAM_KEY = "src/main/resources/ssl/selfsigned.key"
  var SSL_SELF_PASSWORD = "vesoft"
  var CONN_TIMEOUT = 6000
  var CONN_RETRY = 2
  var USER = "root"
  var PASSWORD = "nebula"

  def createNebulaConnConfig(): NebulaConnectionConfig = {
    var config = NebulaConnectionConfig.builder()
      .withMetaAddress(META_ADDRESS)
      .withGraphAddress(GRAPH_ADDRESS)

    if (CONN_RETRY > 0) {
      config = config.withConnectionRetry(CONN_RETRY)
    }
    if (CONN_TIMEOUT > 0) {
      config = config.withTimeout(CONN_TIMEOUT)
    }
    if (META_SSL) {
      config = config.withEnableMetaSSL(true)
    }
    if (GRAPH_SSL) {
      config = config.withEnableGraphSSL(true)
    }
    if (META_SSL && SSL_SIGN_TYPE == "CA") {
      config = config.withSSLSignType(SSLSignType.CA)
        .withCaSSLSignParam(SSL_CA_PARAM_PEM, SSL_CA_PARAM_CRT, SSL_CA_PARAM_KEY)
    }
    if (META_SSL && SSL_SIGN_TYPE == "SELF") {
      config.withSSLSignType(SSLSignType.SELF)
        .withSelfSSLSignParam(SSL_SELF_PARAM_PEM, SSL_SELF_PARAM_KEY, SSL_SELF_PASSWORD)
    }
    config.build()
  }
}
